#include "CubemapToPanorama.h"
#include "main.h"
#include "Color.h"

#include <math.h>
#include <stdio.h>
#include <vector>
#include <algorithm>

class CubemapToPanorama {
public:
	const int SIDE_LEFT = 0;
	const int SIDE_RIGHT = 1;
	const int SIDE_BOTTOM = 2;
	const int SIDE_TOP = 3;
	const int SIDE_FRONT = 4;
	const int SIDE_BACK = 5;

	int width, height;

	godot_color default_color;

	struct Side {
		int width, height;
		const uint8_t *ptr;
	} side[6];

	uint8_t *output;

	float sky_energy;

	float sun_angle_max, sun_angle_min;
	godot_color sun_linear;
	float sun_curve, sun_latitude, sun_longitude;

private:
	bool use_estimated_color;

	struct EstimatedParam {
		float a, b, c;
		float residue;

		float evaluate(float f) const {
			return a + b * powf(f, c);
		}
	};

	EstimatedParam estimated_param[3];

	class EstimateColor {
	public:
		// points
		std::vector<float> xx, yy[3];

	private:
		double n, sum_y[3], sum_y_square[3];

	public:
		void init() {
			n = (double)xx.size();

			for (int idx = 0; idx < 3; idx++) {
				double sy = 0.0, syy = 0.0;
				for (auto y : yy[idx]) {
					sy += y;
					syy += y * y;
				}
				sum_y[idx] = sy;
				sum_y_square[idx] = syy;
			}
		}

		// fix c solve a, b using least square
		void estimate_ab(int idx, EstimatedParam &param) const {
			double sum_xc = 0.0, sum_xc_square = 0.0, sum_xc_y = 0.0;
			for (size_t i = 0; i < xx.size(); i++) {
				float xc = powf(xx[i], param.c);
				sum_xc += xc;
				sum_xc_square += xc * xc;
				sum_xc_y += xc * yy[idx][i];
			}

			double det_inv = 1.0 / (n * sum_xc_square - sum_xc * sum_xc);

			param.a = float((sum_xc_square * sum_y[idx] - sum_xc * sum_xc_y) * det_inv);
			param.b = float((n * sum_xc_y - sum_xc * sum_y[idx]) * det_inv);

			param.residue = float(sum_y_square[idx] - param.a * sum_y[idx] - param.b * sum_xc_y);
		}

		// estimate a,b,c using 0.618 method
		EstimatedParam estimate(int idx) const {
			float cc_begin, cc_end;
			EstimatedParam param[2];

			cc_begin = 0.1f;
			cc_end = 10.0f;
			param[0].c = cc_begin + (cc_end - cc_begin) * 0.381966011f;
			param[1].c = cc_begin + (cc_end - cc_begin) * 0.618033989f;
			estimate_ab(idx, param[0]);
			estimate_ab(idx, param[1]);

			while (cc_end - cc_begin > 1e-3f) {
				if (param[0].residue < param[1].residue) {
					cc_end = param[1].c;
					param[1] = param[0];
					param[0].c = cc_begin + (cc_end - cc_begin) * 0.381966011f;
					estimate_ab(idx, param[0]);
				} else {
					cc_begin = param[0].c;
					param[0] = param[1];
					param[1].c = cc_begin + (cc_end - cc_begin) * 0.618033989f;
					estimate_ab(idx, param[1]);
				}
			}

			if (param[0].residue < param[1].residue) {
				return param[0];
			} else {
				return param[1];
			}
		}
	};

public:
	CubemapToPanorama()
		: width(0), height(0)
		, output(NULL)
		, use_estimated_color(false)
		, sky_energy(1.0f)
		, sun_angle_max(1.6f), sun_angle_min(0.016f)
		, sun_linear(godot_color{ 0.0f, 0.0f, 0.0f, 1.0f })
		, sun_curve(0.05f)
		, sun_latitude(0.6f)
		, sun_longitude(0.0f)
	{
	}

	void estimate_color() {
		// we already have top face
		if (side[SIDE_TOP].ptr != NULL) {
			use_estimated_color = false;
			return;
		}

		// y = a + b * pow(x, c)
		// where x = (angle between vector and +Y) / (PI / 2)
		//       y = color
		EstimateColor est;

		// collect points
		collect_points(side[SIDE_LEFT], est.xx, est.yy[0], est.yy[1], est.yy[2]);
		collect_points(side[SIDE_RIGHT], est.xx, est.yy[0], est.yy[1], est.yy[2]);
		collect_points(side[SIDE_FRONT], est.xx, est.yy[0], est.yy[1], est.yy[2]);
		collect_points(side[SIDE_BACK], est.xx, est.yy[0], est.yy[1], est.yy[2]);

		// too few points
		if (est.xx.size() < 16) {
			use_estimated_color = false;
			return;
		}

		use_estimated_color = true;

		est.init();

		for (int idx = 0; idx < 3; idx++) {
			estimated_param[idx] = est.estimate(idx);

			// DEBUG
			printf("result %d: c=%g,a=%g,b=%g,residue=%g\n", idx, estimated_param[idx].c,
				estimated_param[idx].a, estimated_param[idx].b, estimated_param[idx].residue);
		}
	}

	void calculate() {
		std::vector<float> cos_phi, sin_phi;

		for (int x = 0; x < width; x++) {
			float phi = (x + 0.5f) / width * (2.0f * 3.14159265f);
			cos_phi.push_back(cosf(phi));
			sin_phi.push_back(sinf(phi));
		}

		uint32_t *ptr = (uint32_t*)output;

		godot_color estimated_color = default_color;

		float sun_u, sun_v, sun_w;
		float sun_angle_inv;
		sun_w = cosf(sun_latitude);
		sun_u = sun_w * sinf(sun_longitude);
		sun_w = sun_w * cosf(sun_longitude);
		sun_v = sinf(sun_latitude);
		sun_angle_inv = 1.0f / (sun_angle_max - sun_angle_min);
		bool create_sun = sun_linear.r > 1e-2f || sun_linear.g > 1e-2f || sun_linear.b > 1e-2f;

		for (int y = 0; y < height; y++) {
			float theta = (y + 0.5f) / height * 3.14159265f;
			float v = cosf(theta);
			float sin_theta = sinf(theta);

			if (use_estimated_color && theta < 1.5f) {
				float f = theta * (2.0f / 3.14159265f);
				estimated_color.r = estimated_param[0].evaluate(f);
				estimated_color.g = estimated_param[1].evaluate(f);
				estimated_color.b = estimated_param[2].evaluate(f);
			}

			for (int x = 0; x < width; x++) {
				godot_color clr;
				float u = sin_theta * sin_phi[x], w = sin_theta * cos_phi[x];
				tex_cube(u, v, w, clr, estimated_color);
				clr.r *= sky_energy;
				clr.g *= sky_energy;
				clr.b *= sky_energy;
				if (create_sun) {
					float sun_angle = acosf(Math::clamp(sun_u * u + sun_v * v + sun_w * w, -1.0f, 1.0f));
					if (sun_angle < sun_angle_min) {
						clr = Color::blend(clr, sun_linear);
					} else if (sun_angle < sun_angle_max) {
						float c2 = (sun_angle - sun_angle_min) * sun_angle_inv;
						c2 = Math::ease(c2, sun_curve);
						clr = Color::linear_interpolate(Color::blend(clr, sun_linear), clr, c2);
					}
				}
				*(ptr++) = Color::to_rgbe9995(clr);
			}
		}
	}

private:
	void collect_points(const Side& side, std::vector<float>& xx, std::vector<float>& yr, std::vector<float>& yg, std::vector<float>& yb) {
		const uint8_t *ptr = side.ptr;
		if (ptr == NULL) return;

		for (int y = 0, ym = side.height / 4; y < ym; y++) {
			float fy = 1.0f - 2.0f * (y + 0.5f) / side.height;
			float fy_square_plus_one = fy * fy + 1.0f;

			for (int x = 0; x < side.width; x++) {
				float fx = 2.0f * (x + 0.5f) / side.width - 1.0f;

				xx.push_back(acosf(fy / sqrtf(fx * fx + fy_square_plus_one)) * (2.0f / 3.14159265f));
				yr.push_back(Color::to_linear(*(ptr++)));
				yg.push_back(Color::to_linear(*(ptr++)));
				yb.push_back(Color::to_linear(*(ptr++)));
			}
		}
	}

	inline void tex_integer(const Side& side, int x, int y, godot_color& clr, float f) const {
		if (side.ptr) {
			const uint8_t *ptr = side.ptr + (((intptr_t)y * side.width + x) * 3);
			clr.r += Color::to_linear(ptr[0]) * f;
			clr.g += Color::to_linear(ptr[1]) * f;
			clr.b += Color::to_linear(ptr[2]) * f;
		} else {
			clr.r += default_color.r * f;
			clr.g += default_color.g * f;
			clr.b += default_color.b * f;
		}
	}

	inline void tex(const Side& side, float u, float v, godot_color& clr, float& factor) const {
		u = u * 0.5f * side.width - 0.5f;
		v = v * 0.5f * side.height - 0.5f;

		int x = (int)(u + 2.0f) - 2, y = (int)(v + 2.0f) - 2;

		if (x < -1 || x >= side.width || y < -1 || y >= side.height) return;

		u -= (float)x; v -= (float)y;

		if (y >= 0) {
			if (x >= 0) {
				float f = (1.0f - u) * (1.0f - v);
				tex_integer(side, x, y, clr, f);
				factor += f;
			}
			if (x + 1 < side.width) {
				float f = u * (1.0f - v);
				tex_integer(side, x + 1, y, clr, f);
				factor += f;
			}
		}
		if (y + 1 < side.height) {
			if (x >= 0) {
				float f = (1.0f - u) * v;
				tex_integer(side, x, y + 1, clr, f);
				factor += f;
			}
			if (x + 1 < side.width) {
				float f = u * v;
				tex_integer(side, x + 1, y + 1, clr, f);
				factor += f;
			}
		}
	}

	// NOTE: this is tuned for Ballance skybox files
	inline void tex_cube(float x, float y, float z, godot_color& clr, const godot_color& estimated_color) const {
		clr = godot_color{ 0.0f, 0.0f, 0.0f, 1.0f };
		float factor = 0.0f;

		float abs_x = x < 0.0f ? -x : x;
		float abs_y = y < 0.0f ? -y : y;
		float abs_z = z < 0.0f ? -z : z;

		const float THRESHOLD = 0.55f; // slightly smaller than sqrt(1/3)

		if (abs_x > THRESHOLD) {
			tex(side[x > 0.0f ? SIDE_RIGHT : SIDE_LEFT], 1.0f - z / x, 1.0f - y / abs_x, clr, factor);
		}
		if (abs_z > THRESHOLD) {
			tex(side[z > 0.0f ? SIDE_FRONT : SIDE_BACK], 1.0f + x / z, 1.0f - y / abs_z, clr, factor);
		}
		if (use_estimated_color) {
			if (y < -THRESHOLD) {
				tex(side[SIDE_BOTTOM], 1.0f + x / abs_y, 1.0f + z / y, clr, factor);
			} else if (y > THRESHOLD) {
				float f = std::max(abs_x, abs_z);
				f = y > f ? 1.0f : (y - THRESHOLD) / (f - THRESHOLD);
				factor += (1.0f - factor) * f;
				clr.r += (estimated_color.r - clr.r) * f;
				clr.g += (estimated_color.g - clr.g) * f;
				clr.b += (estimated_color.b - clr.b) * f;
			}
		} else if (abs_y > THRESHOLD) {
			tex(side[y > 0.0f ? SIDE_TOP : SIDE_BOTTOM], 1.0f + x / abs_y, 1.0f + z / y, clr, factor);
		}

		if (factor < 1e-3f) {
			clr = default_color;
		} else {
			factor = 1.0f / factor;
			clr.r *= factor;
			clr.g *= factor;
			clr.b *= factor;
		}
	}
};

// input should be FORMAT_RGB8 sRGB color space with width multiple of 4,
// output should be FORMAT_RGBE9995
godot_variant cubemap_to_panorama(godot_object *p_instance, void *p_method_data,
	void *p_user_data, int p_num_args, godot_variant **p_args)
{
	CubemapToPanorama ctp;

	godot_array arr = gd::api->godot_variant_as_array(p_args[0]);
	godot_int m = gd::api->godot_array_size(&arr);
	godot_variant v;

	godot_pool_byte_array_read_access *pbara[6] = {};

	for (int i = 0; i < 6; i++) {
		v = gd::api->godot_array_get(&arr, i * 3);
		ctp.side[i].width = (int)gd::api->godot_variant_as_int(&v);

		v = gd::api->godot_array_get(&arr, i * 3 + 1);
		ctp.side[i].height = (int)gd::api->godot_variant_as_int(&v);

		ctp.side[i].ptr = NULL;

		if (ctp.side[i].width > 0 && ctp.side[i].height > 0) {
			v = gd::api->godot_array_get(&arr, i * 3 + 2);
			godot_pool_byte_array pba = gd::api->godot_variant_as_pool_byte_array(&v);
			pbara[i] = gd::api->godot_pool_byte_array_read(&pba);
			if (pbara[i]) ctp.side[i].ptr = gd::api->godot_pool_byte_array_read_access_ptr(pbara[i]);
		}
	}

	int max_size = 1;
	for (int i = 0; i < 6; i++) {
		if (max_size < ctp.side[i].width) max_size = ctp.side[i].width;
		if (max_size < ctp.side[i].height) max_size = ctp.side[i].height;
	}
	for (int i = 0; i < 6; i++) {
		if (ctp.side[i].width <= 0) ctp.side[i].width = max_size;
		if (ctp.side[i].height <= 0) ctp.side[i].height = max_size;
	}

	ctp.width = 2048;
	ctp.height = 1024;
	if (m > 18) {
		v = gd::api->godot_array_get(&arr, 18);
		ctp.width = (int)gd::api->godot_variant_as_int(&v);
	}
	if (m > 19) {
		v = gd::api->godot_array_get(&arr, 19);
		ctp.height = (int)gd::api->godot_variant_as_int(&v);
	}

	ctp.default_color = godot_color{ 0.0f, 0.0f, 0.0f, 1.0f };
	if (m > 20) {
		v = gd::api->godot_array_get(&arr, 20);
		godot_color clr = gd::api->godot_variant_as_color(&v);
		ctp.default_color = Color::to_linear(clr);
	}

	bool estimate_color = true;
	if (m > 21) {
		v = gd::api->godot_array_get(&arr, 21);
		estimate_color = gd::api->godot_variant_as_bool(&v);
	}

	if (m > 22) {
		v = gd::api->godot_array_get(&arr, 22);
		ctp.sky_energy = (float)gd::api->godot_variant_as_real(&v);
	}
	if (m > 23) {
		v = gd::api->godot_array_get(&arr, 23);
		ctp.sun_angle_max = (float)gd::api->godot_variant_as_real(&v);
	}
	if (m > 24) {
		v = gd::api->godot_array_get(&arr, 24);
		ctp.sun_angle_min = (float)gd::api->godot_variant_as_real(&v);
	}
	if (m > 25) {
		v = gd::api->godot_array_get(&arr, 25);
		ctp.sun_linear = gd::api->godot_variant_as_color(&v);
	}
	if (m > 26) {
		v = gd::api->godot_array_get(&arr, 26);
		ctp.sun_curve = (float)gd::api->godot_variant_as_real(&v);
	}
	if (m > 27) {
		v = gd::api->godot_array_get(&arr, 27);
		ctp.sun_latitude = (float)gd::api->godot_variant_as_real(&v);
	}
	if (m > 28) {
		v = gd::api->godot_array_get(&arr, 28);
		ctp.sun_longitude = (float)gd::api->godot_variant_as_real(&v);
	}

	godot_pool_byte_array pba;
	gd::api->godot_pool_byte_array_new(&pba);
	gd::api->godot_pool_byte_array_resize(&pba, (godot_int)ctp.width * ctp.height * 4);
	godot_pool_byte_array_write_access *pbawa = gd::api->godot_pool_byte_array_write(&pba);
	ctp.output = gd::api->godot_pool_byte_array_write_access_ptr(pbawa);

	if (estimate_color) ctp.estimate_color();
	ctp.calculate();

	gd::api->godot_pool_byte_array_write_access_destroy(pbawa);

	for (int i = 0; i < 6; i++) {
		if (pbara[i]) gd::api->godot_pool_byte_array_read_access_destroy(pbara[i]);
	}

	godot_variant ret;
	gd::api->godot_variant_new_pool_byte_array(&ret, &pba);
	gd::api->godot_pool_byte_array_destroy(&pba);

	return ret;
}

void cubemap_to_panorama_register_methods(void *p_handle) {
	godot_instance_create_func create = { NULL, NULL, NULL };
	create.create_func = &dummy_constructor;

	godot_instance_destroy_func destroy = { NULL, NULL, NULL };
	destroy.destroy_func = &dummy_destructor;

	gdns::api->godot_nativescript_register_class(p_handle, "NativeCubemapToPanorama", "Reference",
		create, destroy);

	godot_instance_method method = { NULL, NULL, NULL };
	method.method = &cubemap_to_panorama;

	godot_method_attributes attributes = { GODOT_METHOD_RPC_MODE_DISABLED };

	gdns::api->godot_nativescript_register_method(p_handle, "NativeCubemapToPanorama", "cubemap_to_panorama",
		attributes, method);
}
